#include "http_server.h"
#if HTTPD_SERVER_ENABLE_WS

#include "esp_random.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include "http_tools.h"
#include "driver/uart.h"
#include "ota.h"

#include "api_map.h"

#include "esp_log.h"
#define TAG "httpd.wss"
#define log_e(...) ESP_LOGE(TAG, __VA_ARGS__)
#define log_i(...) ESP_LOGI(TAG, __VA_ARGS__)
#define log_w(...) ESP_LOGW(TAG, __VA_ARGS__)

#ifdef CONFIG_SPIRAM
#define _malloc_(size) heap_caps_malloc(size, MALLOC_CAP_SPIRAM)
#else
#define _malloc_(size) malloc(size)
#endif

static void ws_async_send(void *arg)
{
    // if (http_server->ws.send_lock)
    //     xSemaphoreTake(http_server->ws.send_lock, portMAX_DELAY);

    httpd_ws_frame_t ws_pkt = {0};
    memset(&ws_pkt, 0, sizeof(httpd_ws_frame_t));
    ws_pkt.type = http_server->ws.msg_count ? HTTPD_WS_TYPE_BINARY : HTTPD_WS_TYPE_PING;
    ws_pkt.payload = http_server->ws.msg_data;
    ws_pkt.len = http_server->ws.msg_count;

    http_server_ws_clinet_t client = http_server->ws.clients;
    while (client)
    {
        if (ws_pkt.type == HTTPD_WS_TYPE_PING)
        {
            if (client->alive_time > esp_log_timestamp())
                continue;
        }
#if HTTPD_SERVER_ENABLE_AUTH
        if (client->auth && client->shell)
#endif
        {
            if (httpd_ws_send_frame_async(http_server->server, client->fd, &ws_pkt) != ESP_OK)
                break;
        }
        client = client->next;
    }

    // if (http_server->ws.send_lock)
    //     xSemaphoreGive(http_server->ws.send_lock);

    http_server->ws.watting_msg = 0; // clear watting
}

uint16_t ws_msg_read_bytes(uint8_t *data, uint16_t len)
{
    if (http_server)
    {
        if (http_server->ws.msg_len)
        {
            http_server_ws_message_t packet = http_server->ws.msg_list;
            uint16_t read = packet->len - packet->read;
            if (read > len)
                read = len;
            memcpy(data, &packet->data[packet->read], read);
            packet->read += read;
            if (packet->len == packet->read)
            {
                http_server->ws.last_fd = packet->fd;
                http_server->ws.msg_list = http_server->ws.msg_list->next;
                http_server->ws.msg_len--;
                free(packet);
            }
            return read;
        }
    }
    return 0;
}

uint16_t ws_msg_reply(uint8_t *data, uint16_t len)
{
    if (http_server)
    {
        if (http_server->ws.last_fd != -1 || http_server->ws.msg_len)
        {
            http_server_ws_message_t packet = http_server->ws.msg_list;
            int fd = packet->fd;
            if (!packet->read && http_server->ws.last_fd != -1)
                fd = http_server->ws.last_fd;
            httpd_ws_frame_t ws_pkt = {0};
            ws_pkt.type = HTTPD_WS_TYPE_BINARY;
            ws_pkt.payload = data;
            ws_pkt.len = len;
            esp_err_t ret = httpd_ws_send_frame_async(http_server->server, fd, &ws_pkt);
            if (ret != ESP_OK)
                return 0;
            return len;
        }
    }
    return 0;
}

uint16_t ws_send_msg(uint8_t *data, uint16_t len)
{
    if (http_server)
    {
        // log_i("send msg(%d) %s", len, (char *)data);
        // if (ota_default_handle->status == OTA_STATUS_RUNNING)
        //     return len;
        if (http_server->ws.process)
            return len;

        if (http_server->ws.in_process)
        {
            httpd_ws_frame_t ws_pkt = {0};
            memset(&ws_pkt, 0, sizeof(httpd_ws_frame_t));
            ws_pkt.type = HTTPD_WS_TYPE_BINARY;
            ws_pkt.payload = data;
            ws_pkt.len = len;
            httpd_ws_send_frame(http_server->ws.req, &ws_pkt);
        }
        else
        {
            http_server->ws.msg_data = data;
            http_server->ws.msg_count = len;
            http_server->ws.watting_msg = 1; // add watting
            httpd_queue_work(http_server->server, ws_async_send, (void *)HTTPD_WS_TYPE_BINARY);
            // ws_async_send(NULL);

            while (http_server->ws.watting_msg)
                vTaskDelay(pdMS_TO_TICKS(10));
        }
    }
    return len;
}

void ws_reply_status(http_server_t server, httpd_ws_frame_t *frame, char *name, int32_t status)
{
    // log_i("%s reply status %ld", name, status);
    uint16_t len = 0;
    char *str = (char *)server->ws.data;
    len += sprintf(&str[len], "{");
    len += sprintf(&str[len], "\"api\":\"%s\",", name);
    len += sprintf(&str[len], "\"data\":%ld", status);
    len += sprintf(&str[len], "}");
    frame->payload = server->ws.data;
    frame->len = len;
}

void ws_reply_data(http_server_t server, httpd_ws_frame_t *frame, char *name, char *data)
{
    log_i("%s reply data %s", name, data);
    uint16_t len = 0;
    char *str = (char *)server->ws.data;
    len += sprintf(&str[len], "{");
    len += sprintf(&str[len], "\"api\":\"%s\",", name);
    len += sprintf(&str[len], "\"data\":%s", data);
    len += sprintf(&str[len], "}");
    frame->payload = server->ws.data;
    frame->len = len;
}

http_server_ws_clinet_t wss_find_client(httpd_req_t *req)
{
    http_server_t server = httpd_get_global_user_ctx(req->handle);
    http_server_ws_clinet_t client = server->ws.clients;
    int fd = httpd_req_to_sockfd(req);
    while (client)
    {
        if (client->fd == fd)
            return client;
        client = client->next;
    }
    return 0;
}

uint8_t wss_alive_handle(httpd_req_t *req, int fd, uint64_t time)
{
    uint8_t re = 0;
    http_server_t server = httpd_get_global_user_ctx(req->handle);
    http_server_ws_clinet_t client = server->ws.clients;
#if HTTPD_SERVER_ENABLE_AUTH
    if (time)
    {
        for (uint8_t i = 0; i < sizeof(server->ws.auth_list) / sizeof(server->ws.auth_list[0]); i++)
        {
            if (server->ws.auth_list[i])
            {
                if ((time - server->ws.auth_time[i]) / 1000 > 300) /* 5分钟未活动, 会话超时处理 */
                {
                    log_w("time : %llu ms : %llu", server->ws.auth_time[i], time);
                    log_w("This Session(%lu) Timeout, is not alive, close it.", server->ws.auth_list[i]);
                    server->ws.auth_list[i] = 0;
                    server->ws.auth_time[i] = 0;
                }
            }
        }
    }
#endif
    while (client)
    {
        if (!time)
        {
            if (client->fd == fd)
            {
                /* 单独删除此客户端 */
                /* 关闭客户端连接, 释放客户端占用内存 */
                httpd_ws_frame_t ws_pkt = {0};
                ws_pkt.type = HTTPD_WS_TYPE_CLOSE;
                httpd_ws_send_frame_async(req->handle, client->fd, &ws_pkt);
                return 1;
            }
        }
        else if (client->alive_time + server->ws.keep_alive_period_ms <= time)
        {
            log_w("This Client Timeout, is not alive, close it. sockfd:%d", client->fd);
            if (client->fd == fd)
                re = 1;
            const char *data = "{\"code\":-1,\"err\":\"Timeout\"}";
            httpd_ws_frame_t ws_pkt = {0};
            ws_pkt.type = HTTPD_WS_TYPE_BINARY;
            ws_pkt.payload = (uint8_t *)data;
            ws_pkt.len = strlen(data);
            if (httpd_ws_send_frame_async(req->handle, client->fd, &ws_pkt) != ESP_OK)
                break;
            ws_pkt.type = HTTPD_WS_TYPE_CLOSE;
            ws_pkt.payload = 0;
            ws_pkt.len = 0;
            if (httpd_ws_send_frame_async(req->handle, client->fd, &ws_pkt) != ESP_OK)
                break;
        }
        if (client->fd == fd && !re)
        {
            /* 刷新存活时间 */
            client->alive_time = time + server->ws.keep_alive_period_ms;
#if HTTPD_SERVER_ENABLE_AUTH
            if (time)
            {
                for (uint8_t i = 0; i < sizeof(server->ws.auth_list) / sizeof(server->ws.auth_list[0]); i++)
                {
                    if (server->ws.auth_list[i] == client->key)
                    {
                        server->ws.auth_time[i] = time;
                        break;
                    }
                }
            }
#endif
        }
        client = client->next;
    }
    return re;
}

esp_err_t wss_handler(httpd_req_t *req)
{
    http_server_t server = httpd_get_global_user_ctx(req->handle);
    if (!server)
    {
        log_e("Not find http server ctx");
        return ESP_FAIL;
    }
    uint64_t time = esp_log_timestamp();
    int fd = httpd_req_to_sockfd(req);
    /* 检查客户端存活情况 */
    /* 更新时间, 如果返回非零值, 则此客户端已被删除 */
    // if (wss_alive_handle(req, fd, time))
    //     return ESP_OK;
    if (req->method == HTTP_GET)
    {
        http_server_ws_clinet_t ws_clinet = _malloc_(sizeof(http_server_ws_clinet_s));
        if (!ws_clinet)
        {
            log_e("Failed to malloc websocket client buffer.");
            return ESP_FAIL;
        }
        memset(ws_clinet, 0, sizeof(http_server_ws_clinet_s));
        ws_clinet->fd = fd;
#if HTTPD_SERVER_ENABLE_AUTH
        ws_clinet->key = key;
#endif
        ws_clinet->alive_time = time + server->ws.keep_alive_period_ms;
        if (!server->ws.clients)
            server->ws.clients = ws_clinet;
        else
        {
            http_server_ws_clinet_t client = server->ws.clients;
            while (client->next)
                client = client->next;
            client->next = ws_clinet;
        }
        server->ws.length++;
        // log_i("Handshake done, the new connection was opened sock %d, length : %d", fd, server->ws.length);
        return ESP_OK;
    }
    if (http_server->ws.process)
        return ESP_OK;
    httpd_ws_frame_t ws_pkt;
    memset(&ws_pkt, 0, sizeof(httpd_ws_frame_t));
    /* 接收数据长度 */
    esp_err_t ret = httpd_ws_recv_frame(req, &ws_pkt, 0);
    if (ret != ESP_OK)
    {
        log_e("httpd_ws_recv_frame failed to get frame len with %d", ret);
        return ret;
    }
    if (ws_pkt.len)
    {
        ws_pkt.payload = server->ws.data;
        /* 接收数据 */
        ret = httpd_ws_recv_frame(req, &ws_pkt, MIN(ws_pkt.len, 4096));
        if (ret != ESP_OK)
        {
            log_e("httpd_ws_recv_frame failed with %d", ret);
            return ret;
        }
    }
    // log_i("frame len is %d, type : %d, sock : %d", ws_pkt.len, ws_pkt.type, httpd_req_to_sockfd(req));
    /* 如果是接收到pong消息 */
    // if (ws_pkt.type == HTTPD_WS_TYPE_PONG)
    //     log_i("Received PONG message, fd:%d", fd);
    // else
    /* 接收到ping消息 */
    if (ws_pkt.type == HTTPD_WS_TYPE_PING)
    {
        // log_i("Got a WS PING frame, Replying PONG");
        ws_pkt.type = HTTPD_WS_TYPE_PONG;
        ret = httpd_ws_send_frame(req, &ws_pkt);
        if (ret != ESP_OK)
            log_e("httpd_ws_send_frame failed with %d", ret);
        return ret;
    }
    /* 接收到连接断开消息 */
    else if (ws_pkt.type == HTTPD_WS_TYPE_CLOSE)
    {
        log_i("Got a WS Close frame, delete it.");
        /* 删除客户端 */
        wss_alive_handle(req, fd, 0);
        return ESP_OK;
    }
    else if (ws_pkt.type == HTTPD_WS_TYPE_TEXT || ws_pkt.type == HTTPD_WS_TYPE_BINARY)
    {
        http_server_ws_clinet_t client = wss_find_client(req);
        if (!client)
        {
            ws_pkt.type = HTTPD_WS_TYPE_CLOSE;
            ws_pkt.len = 0;
            ws_pkt.payload = 0;
            httpd_ws_send_frame(req, &ws_pkt);
            return ESP_OK;
        }
        server->ws.in_process = 1;
        server->ws.req = req;
        cJSON *json = NULL;
        /* 接收到文本消息 */
        if (ws_pkt.type == HTTPD_WS_TYPE_TEXT)
        {
            if (ws_pkt.len < sizeof(server->ws.data))
                ws_pkt.payload[ws_pkt.len] = '\0';
            if (!strnstr((const char *)ws_pkt.payload, "ping", ws_pkt.len))
                // log_i("Received packet(%u) with message(%p): '%s'", ws_pkt.len, ws_pkt.payload, ws_pkt.payload);
                log_i("%s", ws_pkt.payload);
            ws_pkt.final = false;
            ws_pkt.fragmented = false;
            ws_pkt.type = HTTPD_WS_TYPE_BINARY;
            json = cJSON_Parse((char *)ws_pkt.payload);
            ws_pkt.len = 0;
            if (!cJSON_IsObject(json))
            {
                ws_pkt.payload = (uint8_t *)"{\"code\":-1, \"err\":\"The data is not a JSON obj object\"}";
                ws_pkt.len = strlen((char *)ws_pkt.payload);
                goto ws_msg_end;
            }
            if (!cJSON_HasObjectItem(json, "api"))
            {
                ws_pkt.payload = (uint8_t *)"{\"code\":-1, \"err\":\"Missing parameter API\"}";
                ws_pkt.len = strlen((char *)ws_pkt.payload);
                goto ws_msg_end;
            }
            char *str = cJSON_GetStringValue(cJSON_GetObjectItem(json, "api"));
            if (!str)
            {
                ws_pkt.payload = (uint8_t *)"{\"code\":-1, \"err\":\"API parameter Value Type is not a string\"}";
                ws_pkt.len = strlen((char *)ws_pkt.payload);
                goto ws_msg_end;
            }
            if (!strcmp(str, "ping"))
            {
                ws_pkt.payload = (uint8_t *)"{\"api\":\"pong\"}";
                ws_pkt.len = strlen((char *)ws_pkt.payload);
                goto ws_msg_end;
            }

            cJSON *data = cJSON_GetObjectItem(json, "data");
            memset(server->ws.data, 0, sizeof(server->ws.data));
            server->cur_req = req;
            api_call(str, server, &ws_pkt, data);
        }
        /* 二进制内容 */
        else if (ws_pkt.type == HTTPD_WS_TYPE_BINARY)
        {
            if (server->ws.file)
            {
                size_t write = fwrite(ws_pkt.payload, 1, ws_pkt.len, server->ws.file);
                if (write != ws_pkt.len)
                    ws_reply_status(server, &ws_pkt, "change_file", -1);
                else
                {
                    server->ws.file_recv_size += write;
                    if (server->ws.file_recv_size >= server->ws.file_size)
                    {
                        fclose(server->ws.file);
                        server->ws.file = NULL;
                        ws_reply_status(server, &ws_pkt, "change_file", 1);
                    }
                }
                log_i("ws binary write p:%d w:%d r:%ld t:%ld", ws_pkt.len, write, server->ws.file_recv_size, server->ws.file_size);
            }
        }
    ws_msg_end:
        if (json)
            cJSON_Delete(json);
        /* 回传消息 */
        if (ws_pkt.len)
        {
            // if (ws_pkt.len < 255)
            //     log_i("ws send %s %d", ws_pkt.payload, ws_pkt.len);
            ret = httpd_ws_send_frame(req, &ws_pkt);
            if (ret != ESP_OK)
                log_e("httpd_ws_send_frame failed with %d", ret);
        }
        server->ws.in_process = 0;
        return ret;
    }
    return ESP_OK;
}
#endif
